package tagdb.model;

import tagdb.algorithm.match.PeptideSpectrumMatch;

import java.util.*;

/**
 * <b>TagSearchResult</b>
 * <p>
 * This class holds all information about retrieved hits from the database search.
 * </p>
 *
 * @author Thilo Muth
 */
public class TagSearchResult implements Result {
    /**
     * The protein family hits.
     */
    private Map<String, ProteinFamilyHit> hits = new HashMap<String, ProteinFamilyHit>();

    /**
     * Get the list of all peptide spectrum matches.
     */
    private List<PeptideSpectrumMatch> psms = new ArrayList<PeptideSpectrumMatch>();

    /**
     * The file name.
     */
    private String filename;

    /**
     * Default constructor for search result object initialization.
     */
    public TagSearchResult(String filename) {
        this.filename = filename;
    }

    /**
     * Adding a protein to the protein hit set.
     *
     * @param proteinHit The ProteinHit
     */
    public void addProtein(ProteinHit proteinHit) {

        String accession = proteinHit.getAccession();
        String proteinFamily = proteinHit.getFamily();

        if (hits.containsKey(proteinFamily)) {
            // Current family hit
            ProteinFamilyHit currentFamilyHit = hits.get(proteinFamily);
            TreeMap<String, ProteinHit> proteinHits = currentFamilyHit.getProteinHits();

            // Check if protein hit is already in the protein hit set.
            if (proteinHits.containsKey(accession)) {
                // Current protein hit
                ProteinHit currentProteinHit = proteinHits.get(accession);

                // Get the first - and only - peptide hit.
                PeptideHit peptideHit = proteinHit.getSinglePeptideHit();

                // Get the current peptide hits.
                TreeMap<String, PeptideHit> currentPeptideHits = currentProteinHit.getPeptideHits();

                // Check if peptide hit is not already in the protein hit
                if (!currentPeptideHits.containsKey(peptideHit.getSequence())) {
                    // Add the peptide hit to the protein hit.
                    currentPeptideHits.put(peptideHit.getSequence(), peptideHit);
                    currentProteinHit.setPeptideHits(currentPeptideHits);
                } else { // If peptide hit is already in the result check the actual PSM.
                    // Returns the single PSM.
                    PeptideSpectrumMatch psm = (PeptideSpectrumMatch) peptideHit.getSingleSpectrumMatch();
                    PeptideHit currentPeptideHit = currentPeptideHits.get(peptideHit.getSequence());

                    TreeMap<String, PeptideSpectrumMatch> currentPsms = currentPeptideHit.getSpectrumMatches();

                    // Current PSM
                    if (!currentPsms.containsKey(psm.getSpectrumTitle())) {
                        currentPsms.put(psm.getSpectrumTitle(), psm);
                        currentPeptideHit.setSpectrumMatches(currentPsms);
                        currentPeptideHits.put(currentPeptideHit.getSequence(), currentPeptideHit);
                        currentProteinHit.setPeptideHits(currentPeptideHits);
                    } else {
                        PeptideSpectrumMatch currentPsm = currentPsms.get(psm.getSpectrumTitle());
                        currentPsms.put(psm.getSpectrumTitle(), currentPsm);
                        currentPeptideHit.setSpectrumMatches(currentPsms);
                        currentPeptideHits.put(currentPeptideHit.getSequence(), currentPeptideHit);
                        currentProteinHit.setPeptideHits(currentPeptideHits);
                    }
                }
                proteinHits.put(accession, currentProteinHit);
            } else {
                proteinHits.put(accession, proteinHit);
            }
            currentFamilyHit.setProteinHits(proteinHits);
            hits.put(proteinFamily, currentFamilyHit);
        } else {
            String familyDesc = proteinHit.getDescription().substring(0, proteinHit.getDescription().indexOf("OS"));
            hits.put(proteinFamily, new ProteinFamilyHit(proteinHit.getFamily(), familyDesc, proteinHit));
        }

    }

    /**
     * Returns the protein hit for a particular accession.
     *
     * @param familyName The protein family name.
     * @return The ProteinHit object.
     */
    public ProteinFamilyHit getProteinFamilyHit(String familyName) {
        return hits.get(familyName);
    }

    /**
     * Returns the protein family hits.
     *
     * @return The protein family hits.
     */
    public Map<String, ProteinFamilyHit> getResultMap() {
        return hits;
    }

    /**
     * Returns the filename.
     * @return The filename
     */
    public String getFilename() {
        return filename;
    }

    /**
     * Returns the list of the actual peptide-spectrum-matches.
     * @return The list of the actual peptide-spectrum-matches.
     */
    public List<PeptideSpectrumMatch> getPeptideSpectrumMatches() {
        return psms;
    }

    /**
     * Adds a peptide-spectrum-match.
     * @param psm PeptideSpectrumMatch
     */
    public void addPeptideSpectrumMatch(PeptideSpectrumMatch psm) {
        this.psms.add(psm);
    }
}
